{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build an Order Completion Agent with an Artifact Editor Tool\n",
    "\n",
    "In this example, we'll build a chat assistant that is designed to fill in a custom 'form'.\n",
    "\n",
    "As an example use-case, we'll build an order taking assistant that will need to get a few set pieces of information from the end-user before proceeding. Like their delivery address, and the contents of their order.\n",
    "\n",
    "To build this, we're using the new `ArtifactEditorToolSpec` and `ArtifactMemoryBlock`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index llama-index-tools-artifact-editor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from getpass import getpass\n",
    "from pydantic import BaseModel, Field\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core.memory import Memory\n",
    "from llama_index.core.agent.workflow import (\n",
    "    FunctionAgent,\n",
    "    AgentWorkflow,\n",
    "    ToolCallResult,\n",
    "    AgentStream,\n",
    ")\n",
    "from llama_index.tools.artifact_editor import (\n",
    "    ArtifactEditorToolSpec,\n",
    "    ArtifactMemoryBlock,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"OpenAI API Key: \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Pizza(BaseModel):\n",
    "    name: str = Field(description=\"The name of the pizza\")\n",
    "    remove: list[str] | None = Field(\n",
    "        description=\"If exists, the ingredients the customer requests to remove\",\n",
    "        default=None,\n",
    "    )\n",
    "    add: list[str] | None = Field(\n",
    "        description=\"If exists, the ingredients the customer requests to be added\",\n",
    "        default=None,\n",
    "    )\n",
    "\n",
    "\n",
    "class Address(BaseModel):\n",
    "    street_address: str = Field(\n",
    "        description=\"The street address of the customer\"\n",
    "    )\n",
    "    city: str = Field(description=\"The city of the customer\")\n",
    "    state: str = Field(description=\"The state of the customer\")\n",
    "    zip_code: str = Field(description=\"The zip code of the customer\")\n",
    "\n",
    "\n",
    "class Order(BaseModel):\n",
    "    pizzas: list[Pizza] | None = Field(\n",
    "        description=\"The pizzas ordered by the customer\", default=None\n",
    "    )\n",
    "    address: Address | None = Field(\n",
    "        description=\"The full address of the customer\", default=None\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tool_spec = ArtifactEditorToolSpec(Order)\n",
    "tools = tool_spec.to_tool_list()\n",
    "\n",
    "# Initialize the memory\n",
    "memory = Memory.from_defaults(\n",
    "    session_id=\"order_editor\",\n",
    "    memory_blocks=[ArtifactMemoryBlock(artifact_spec=tool_spec)],\n",
    "    token_limit=60000,\n",
    "    chat_history_token_ratio=0.7,\n",
    ")\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4.1\")\n",
    "\n",
    "agent = AgentWorkflow(\n",
    "    agents=[\n",
    "        FunctionAgent(\n",
    "            llm=llm,\n",
    "            tools=tools,\n",
    "            system_prompt=\"\"\"You are a worker at a Pizzeria. Your job is to talk to users and gather order information.\n",
    "            At every step, you should check the order completeness before responding to the user, and ask for any possibly\n",
    "            missing information.\"\"\",\n",
    "        )\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def chat():\n",
    "    while True:\n",
    "        user_msg = input(\"User: \").strip()\n",
    "        if user_msg.lower() in [\"exit\", \"quit\"]:\n",
    "            print(\"\\n------ORDER COMPLETION-------\\n\")\n",
    "            print(\n",
    "                f\"The Order was placed with the following Order schema:\\n: {json.dumps(tool_spec.get_current_artifact(), indent=4)}\"\n",
    "            )\n",
    "            break\n",
    "\n",
    "        handler = agent.run(user_msg, memory=memory)\n",
    "        async for ev in handler.stream_events():\n",
    "            if isinstance(ev, AgentStream):\n",
    "                print(ev.delta, end=\"\", flush=True)\n",
    "            elif isinstance(ev, ToolCallResult):\n",
    "                print(\n",
    "                    f\"\\n\\nCalling tool: {ev.tool_name} with kwargs: {ev.tool_kwargs}\"\n",
    "                )\n",
    "\n",
    "        # response = await handler\n",
    "        # print(str(response))\n",
    "        print(\"\\n\\nCurrent artifact: \", tool_spec.get_current_artifact())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello! Welcome to our pizzeria. Would you like to place an order? If so, could you please tell me what kind of pizza you’d like?\n",
      "\n",
      "Current artifact:  None\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Calling tool: create_artifact with kwargs: {'pizzas': [{'name': 'pepperoni', 'add': ['olives']}, {'name': 'margherita'}]}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You’ve ordered:\n",
      "- 1 Pepperoni pizza with added olives\n",
      "- 1 Margherita pizza\n",
      "\n",
      "To complete your order, could you please provide your delivery address (street address, city, state, and zip code)?\n",
      "\n",
      "Current artifact:  {'pizzas': [{'name': 'pepperoni', 'remove': None, 'add': ['olives']}, {'name': 'margherita', 'remove': None, 'add': None}], 'address': None}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Calling tool: apply_patch with kwargs: {'patch': {'operations': [{'op': 'replace', 'path': '/address', 'value': {'street_address': '1 Sesame Street', 'city': 'Amsterdam', 'state': 'North-Holand', 'zip_code': '1111AB'}}]}}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Thank you! Your order is now complete:\n",
      "\n",
      "- 1 Pepperoni pizza with added olives\n",
      "- 1 Margherita pizza\n",
      "\n",
      "Delivery address:\n",
      "1 Sesame Street, Amsterdam, North-Holand, 1111AB\n",
      "\n",
      "Would you like to add anything else to your order, or should I proceed with placing it?\n",
      "\n",
      "Current artifact:  {'pizzas': [{'name': 'pepperoni', 'remove': None, 'add': ['olives']}, {'name': 'margherita', 'remove': None, 'add': None}], 'address': {'street_address': '1 Sesame Street', 'city': 'Amsterdam', 'state': 'North-Holand', 'zip_code': '1111AB'}}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Great! Your order has been placed:\n",
      "\n",
      "- 1 Pepperoni pizza with added olives\n",
      "- 1 Margherita pizza\n",
      "\n",
      "Delivery to: 1 Sesame Street, Amsterdam, North-Holand, 1111AB\n",
      "\n",
      "Thank you for ordering with us! Your pizzas will be delivered soon. Have a delicious day!\n",
      "\n",
      "Current artifact:  {'pizzas': [{'name': 'pepperoni', 'remove': None, 'add': ['olives']}, {'name': 'margherita', 'remove': None, 'add': None}], 'address': {'street_address': '1 Sesame Street', 'city': 'Amsterdam', 'state': 'North-Holand', 'zip_code': '1111AB'}}\n",
      "\n",
      "------ORDER COMPLETION-------\n",
      "\n",
      "The Order was placed with the following Order schema:\n",
      ": {\n",
      "    \"pizzas\": [\n",
      "        {\n",
      "            \"name\": \"pepperoni\",\n",
      "            \"remove\": null,\n",
      "            \"add\": [\n",
      "                \"olives\"\n",
      "            ]\n",
      "        },\n",
      "        {\n",
      "            \"name\": \"margherita\",\n",
      "            \"remove\": null,\n",
      "            \"add\": null\n",
      "        }\n",
      "    ],\n",
      "    \"address\": {\n",
      "        \"street_address\": \"1 Sesame Street\",\n",
      "        \"city\": \"Amsterdam\",\n",
      "        \"state\": \"North-Holand\",\n",
      "        \"zip_code\": \"1111AB\"\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "await chat()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
